import argparse
import glob
import os
from importlib.resources import files
import sqlite3
from typing import Optional,Dict,Any
import numpy as np
import pandas as pd
import matplotlib as plt
import copy
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from sklearn.preprocessing import PolynomialFeatures
from abc import abstractmethod,ABC
import asyncio

from mixserve import SamplingParams
from mixserve.profiler import Profiler
from mixserve.config import (
    ModelConfig,
    DisaggParallelConfig,
    ParallelConfig,
    CacheConfig,
    ContextStageSchedConfig,
    DecodingStageSchedConfig
)


parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, help='The model to use', default='/data0/model/Qwen3-8B/Qwen3_8B-bin/')
args = parser.parse_args()

package_root = files('mixserve')
db_path=os.path.join(package_root,'benchdb')
file = glob.glob(os.path.join(db_path, "*"))
for f in file:
    if os.path.isfile(f):
        os.remove(f)
profiler= Profiler(
                model_config=ModelConfig(
                model=args.model,
                tokenizer=None
            ),
            dis_para_config=DisaggParallelConfig(
                context=ParallelConfig(
                    tensor_parallel_size=4,
                    pipeline_parallel_size=1
                ),decoding=ParallelConfig(
                tensor_parallel_size=4,
                pipeline_parallel_size=1
            )
        ),
            cache_config=CacheConfig(
                block_size=16,
                max_num_blocks_per_req=256,
                gpu_memory_utilization=0.9,
                cpu_swap_space=1.0
            ),
            profile_config=ContextStageSchedConfig(
                policy="fcfs",
                max_batch_size=64,
                max_tokens_per_batch=16384
            ),
                decode_profile_config=DecodingStageSchedConfig(
                policy="fcfs",
                max_batch_size=64,
                max_tokens_per_batch=16384
            )
            
            )

asyncio.run(profiler.profile())
